import math
import random


def kmeans_plus_plus_init(data, k):
    centroids = []

    # 1. 随机选择第一个聚类中心
    centroids.append(random.choice(data))

    # 2. 选择剩余的 k-1 个聚类中心
    for _ in range(1, k):
        # 计算每个数据点到最近聚类中心的距离
        distances = []
        for point in data:
            min_distance = float('inf')
            for centroid in centroids:
                distance = 0
                for j in range(len(point)):
                    distance += (point[j] - centroid[j]) ** 2
                min_distance = min(min_distance, distance)
            distances.append(min_distance)

        # 将距离转换为概率（距离越大，被选中的概率越大）
        total_distance = sum(distances)
        probabilities = [dist / total_distance for dist in distances]

        # 根据概率分布选择下一个聚类中心
        next_centroid_idx = random.choices(range(len(data)), weights=probabilities)[0]
        centroids.append(data[next_centroid_idx])

    return centroids


def assign_cluster(x, centroids):
    min_distance = float('inf')
    closest_centroid = 0

    for i, centroid in enumerate(centroids):
        # 计算欧几里得距离
        distance = 0
        for j in range(len(x)):
            distance += (x[j] - centroid[j]) ** 2
        distance = math.sqrt(distance)

        if distance < min_distance:
            min_distance = distance
            closest_centroid = i

    return closest_centroid


def Kmeans(data, k, epsilon=1e-4, max_iterations=100, init_method='random'):
    # 1. 初始化聚类中心
    if init_method == 'kmeans++':
        centroids = kmeans_plus_plus_init(data, k)
        print("使用K-means++初始化")
    else:
        centroids = random.sample(data, k)
        print("使用随机初始化")

    print("初始聚类中心:")
    for i, centroid in enumerate(centroids):
        print(f"  聚类 {i}: {[round(c, 3) for c in centroid]}")

    for iteration in range(max_iterations):
        # 2. 分配数据点到聚类
        clusters = [[] for _ in range(k)]
        assignments = []

        for point in data:
            cluster_idx = assign_cluster(point, centroids)
            clusters[cluster_idx].append(point)
            assignments.append(cluster_idx)

        # 3. 更新聚类中心
        new_centroids = []
        for cluster_points in clusters:
            if not cluster_points:
                # 如果聚类为空，随机选择一个数据点作为中心
                new_centroids.append(random.choice(data))
            else:
                # 计算每个维度的平均值
                new_centroid = []
                for dim in range(len(data[0])):
                    dim_avg = sum(point[dim] for point in cluster_points) / len(cluster_points)
                    new_centroid.append(dim_avg)
                new_centroids.append(new_centroid)

        # 4. 检查收敛
        converged = True
        for i in range(k):
            # 计算新旧聚类中心的距离
            distance = 0
            for dim in range(len(centroids[i])):
                distance += (centroids[i][dim] - new_centroids[i][dim]) ** 2
            distance = math.sqrt(distance)

            if distance > epsilon:
                converged = False
                break

        centroids = new_centroids

        if converged:
            print(f"在第 {iteration + 1} 次迭代后收敛")
            break
    else:
        print(f"达到最大迭代次数 {max_iterations}")

    return centroids, clusters, assignments
if __name__ == "__main__":
    # 创建简单的测试数据
    random.seed(42)

    # 生成三组聚类数据
    test_data = []
    # 第一组
    for _ in range(20):
        test_data.append([random.gauss(2, 0.3), random.gauss(2, 0.3)])
    # 第二组
    for _ in range(20):
        test_data.append([random.gauss(8, 0.3), random.gauss(8, 0.3)])
    # 第三组
    for _ in range(20):
        test_data.append([random.gauss(5, 0.3), random.gauss(2, 0.3)])

    print("数据集大小:", len(test_data))

    # 比较两种初始化方法
    print("\n" + "=" * 50)
    print("使用随机初始化:")
    print("=" * 50)
    centroids_random, clusters_random, assignments_random = Kmeans(
        test_data, k=3, epsilon=0.001, init_method='random'
    )

    print("\n" + "=" * 50)
    print("使用K-means++初始化:")
    print("=" * 50)
    centroids_plus, clusters_plus, assignments_plus = Kmeans(
        test_data, k=3, epsilon=0.001, init_method='kmeans++'
    )

    # 输出结果比较
    print("\n" + "=" * 50)
    print("结果比较:")
    print("=" * 50)

    print("\n随机初始化结果:")
    for i, centroid in enumerate(centroids_random):
        print(f"聚类 {i}: {[round(c, 3) for c in centroid]}, 点数: {len(clusters_random[i])}")

    print("\nK-means++初始化结果:")
    for i, centroid in enumerate(centroids_plus):
        print(f"聚类 {i}: {[round(c, 3) for c in centroid]}, 点数: {len(clusters_plus[i])}")


    # 计算聚类质量（簇内距离和）
    def calculate_wcss(centroids, clusters):
        """计算簇内平方和（Within-Cluster Sum of Squares）"""
        wcss = 0
        for i, cluster_points in enumerate(clusters):
            centroid = centroids[i]
            for point in cluster_points:
                distance = 0
                for dim in range(len(point)):
                    distance += (point[dim] - centroid[dim]) ** 2
                wcss += distance
        return wcss


    wcss_random = calculate_wcss(centroids_random, clusters_random)
    wcss_plus = calculate_wcss(centroids_plus, clusters_plus)

    print(f"\n聚类质量比较:")
    print(f"随机初始化 WCSS: {wcss_random:.4f}")
    print(f"K-means++初始化 WCSS: {wcss_plus:.4f}")
    print(f"改进: {(wcss_random - wcss_plus) / wcss_random * 100:.2f}%")